import torch
import torch.nn as nn
import torch.nn.functional as F
from src.TiDE.TiDE import ResidualBlock
from src.TCN.TCN import *


class QNetWithTCN(nn.Module):
    def __init__(self, state_width=18, action_size=2, dropout=0.05):
        super(QNetWithTCN, self).__init__()
        self.input = nn.Sequential(
            ResidualBlock(state_width, 16, dropout=dropout)
        )
        self.tcn = TemporalConvNet(16, [15, 15, 15, 15], 3, 0.2)
        self.output = nn.Sequential(
            nn.Linear(15, 16),
            nn.LeakyReLU(),
            nn.Dropout(dropout)
        )
        self.value = nn.Sequential(
            nn.Linear(16, 1)
        )
        self.advantage = nn.Sequential(
            nn.Linear(16, action_size),
        )
        # self.softmax = nn.Softmax(dim=1)

    def forward(self, x, position):
        x = self.input(x)
        x = x.permute(0, 2, 1).contiguous()
        x = self.tcn(x)

        x = x[:, :, -1]

        x = self.output(x)
        val = self.value(x)
        adv = self.advantage(x)
        x = val + adv - torch.mean(adv)

        return x
